import sys
import os
import time
import math
import queue

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision

from scene.deformations import Embedder
from utils.general_utils import get_expon_lr_func
from octree_cpp import query_grid

class OctreeDecoder(nn.Module):
    def __init__(self, D=2, W=256, feature_ch=64):
        super(OctreeDecoder, self).__init__()
        self.D = D
        self.W = W
        self.feature_ch = feature_ch

        self.is_time_emb = True
        self.time_pe = 6
        self.is_time_aug = True
        self.time_aug_dim = 30
        self.is_viewpoint_emb = True
        self.viewpoint_pe = 10

        self.tim_embedder = Embedder(input_dim=1, num_freqs=self.time_pe) if self.is_time_emb else None
        self.viewpoint_embedder = Embedder(input_dim=3, num_freqs=self.viewpoint_pe) if self.is_viewpoint_emb else None
        self.tim_aug =  nn.Sequential(
            nn.Linear(self.tim_embedder.out_dim if self.is_time_emb else 1, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, self.time_aug_dim),
        ) if self.is_time_aug else None
        self.time_out_dim = self.time_aug_dim if self.is_time_aug else self.tim_embedder.out_dim if self.is_time_emb else 1
        self.viewpoint_out_dim = self.viewpoint_embedder.out_dim if self.is_viewpoint_emb else 3

        self.net = nn.ModuleList([nn.Linear(self.feature_ch + self.viewpoint_out_dim + self.time_out_dim + 16, self.W)] +
            [nn.Linear(self.W, self.W) for i in range(self.D - 1)])
        self.decoder = nn.Linear(self.W, 3)
        
    def forward(self, x, x_f, v, t):
        t = self.tim_embedder.embed(t) if self.is_time_emb else t
        t = self.tim_aug(t) if self.is_time_aug else t

        v = self.viewpoint_embedder.embed(v) if self.is_viewpoint_emb else v

        h = torch.cat([x, x_f, v, t], dim=-1)
        h = self.net[0](h)
        h = F.relu(h)
        for i, linear in enumerate(self.net):
            if i > 0:
                h = linear(h)
                h = F.relu(h)
        
        rgb = self.decoder(h)
        return rgb



class Octree:
    def __default__parameters(self):
        # Parameters
        self.max_depth = 16
        self.max_points = 64
        self.octree_feature_dim = 64
        self.spatial_lr_scale = 5

        self.decoder = OctreeDecoder(D=2, W=256, \
            feature_ch=self.octree_feature_dim)
        self.decoder.cuda()

        self.max_sh_degree = 3

    def __default__init__(self):
        self.xyz_mins = None
        self.xyz_maxs = None
        self.sons = None
        self.is_leaf_node = None
        self.root = 0
        self.features = None

        self.__default__parameters()

    def __init__(self, points, xyz_min, xyz_max, feature_dc):
        if points is None:
            self.__default__init__()
            return

        self.xyz_mins = [xyz_min.cuda().detach().float()]
        self.xyz_maxs = [xyz_max.cuda().detach().float()]
        self.sons = [torch.zeros(8).cuda().long()]
        self.is_leaf_node = [True]
        self.root = 0

        self.__default__parameters()

        self.grid_fused_color = [torch.zeros(3).cuda().float()]
        self.create_octree(points, feature_dc)
        self.grid_fused_color = torch.stack(self.grid_fused_color, dim=0)

        self.features = torch.rand(len(self.xyz_mins), self.octree_feature_dim).cuda().float().contiguous().requires_grad_(True)
        self.sh_features_dc = self.grid_fused_color.unsqueeze(-1).cuda().transpose(1, 2).contiguous().requires_grad_(True)
        self.sh_features_rest = torch.zeros(self.grid_fused_color.shape[0], 3, (self.max_sh_degree + 1)**2 - 1).cuda().transpose(1, 2).contiguous().requires_grad_(True)
        self.xyz_mins = torch.stack(self.xyz_mins, dim=0)
        self.xyz_maxs = torch.stack(self.xyz_maxs, dim=0)
        self.sons = torch.stack(self.sons, dim=0)
        self.is_leaf_node = torch.tensor(self.is_leaf_node).cuda()

    def cuda(self):
        self.decoder.cuda()        

    def create_octree(self, points, feature_dc):
        q = queue.Queue()
        q.put((self.root, 1))

        # st = time.time()
        while not q.empty():
            node_id, depth = q.get()
            points_filter_mask = ((points >= self.xyz_mins[node_id]) & (points <= self.xyz_maxs[node_id])).all(dim=-1)
            points_filter = points[points_filter_mask]

            if depth >= self.max_depth or points_filter.shape[0] <= self.max_points:
                self.grid_fused_color[node_id] = torch.mean(feature_dc[points_filter_mask].detach(), dim=0) if points_filter.shape[0] != 0 else \
                                                torch.zeros(3).cuda().float()
                continue
            
            self.is_leaf_node[node_id] = False
            sub_xyz_mins, sub_xyz_maxs = self.get_sub_nodes(self.xyz_mins[node_id], self.xyz_maxs[node_id])
            
            for i, (sub_xyz_min, sub_xyz_max) in enumerate(zip(sub_xyz_mins, sub_xyz_maxs)):
                self.xyz_mins.append(sub_xyz_min)
                self.xyz_maxs.append(sub_xyz_max)
                self.sons.append(torch.zeros(8).cuda().long())
                self.sons[node_id][i] = len(self.xyz_mins) - 1
                self.is_leaf_node.append(True)

                self.grid_fused_color.append(torch.zeros(3).cuda().float())
                
                q.put((len(self.xyz_mins) - 1, depth + 1))
            # print('time3', time.time() - st)
            # exit(0)
        
        print('Create octree ending. The number of nodes is {}'.format(len(self.xyz_mins)))

    def get_sub_nodes(self, xyz_min, xyz_max):
        x0, y0, z0 = tuple(xyz_min.cpu().numpy())
        x2, y2, z2 = tuple(xyz_max.cpu().numpy())
        x1, y1, z1 = (x0 + x2) / 2, (y0 + y2) / 2, (z0 + z2) / 2

        return [
            torch.tensor([x0, y0, z0]).cuda().float(),
            torch.tensor([x1, y0, z0]).cuda().float(),
            torch.tensor([x0, y1, z0]).cuda().float(),
            torch.tensor([x0, y0, z1]).cuda().float(),
            torch.tensor([x1, y1, z0]).cuda().float(),
            torch.tensor([x1, y0, z1]).cuda().float(),
            torch.tensor([x0, y1, z1]).cuda().float(),
            torch.tensor([x1, y1, z1]).cuda().float(),
        ], [
            torch.tensor([x1, y1, z1]).cuda().float(),
            torch.tensor([x2, y1, z1]).cuda().float(),
            torch.tensor([x1, y2, z1]).cuda().float(),
            torch.tensor([x1, y1, z2]).cuda().float(),
            torch.tensor([x2, y2, z1]).cuda().float(),
            torch.tensor([x2, y1, z2]).cuda().float(),
            torch.tensor([x1, y2, z2]).cuda().float(),
            torch.tensor([x2, y2, z2]).cuda().float(),
        ]

    def training_setup(self, training_args):
        l = [
            {"params": [self.features], "lr": training_args.grid_lr_init * self.spatial_lr_scale, "name": "octree_features"},
            {"params": list(self.decoder.parameters()), "lr": training_args.deformation_lr_init * self.spatial_lr_scale, "name": "octree_decoder"},
            {"params": [self.sh_features_dc], "lr": training_args.feature_lr, "name": "sh_features_dc"},
            {"params": [self.sh_features_rest], "lr": training_args.feature_lr / 20.0, "name": "sh_features_rest"},
        ]
        self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)

        self.features_scheduler_args = get_expon_lr_func(lr_init=training_args.deformation_lr_init*self.spatial_lr_scale,
                                                    lr_final=training_args.deformation_lr_final*self.spatial_lr_scale,
                                                    lr_delay_mult=training_args.deformation_lr_delay_mult,
                                                    max_steps=training_args.position_lr_max_steps)    
        self.decoder_scheduler_args = get_expon_lr_func(lr_init=training_args.grid_lr_init*self.spatial_lr_scale,
                                                    lr_final=training_args.grid_lr_final*self.spatial_lr_scale,
                                                    lr_delay_mult=training_args.deformation_lr_delay_mult,
                                                    max_steps=training_args.position_lr_max_steps)

    def update_learning_rate(self, iteration):
        for param_group in self.optimizer.param_groups:
            if param_group["name"] == "octree_features":
                lr = self.features_scheduler_args(iteration)
                param_group['lr'] = lr
                # return lr
            if  param_group["name"] == "octree_decoder":
                lr = self.decoder_scheduler_args(iteration)
                param_group['lr'] = lr
                # return lr
        
    def query_rgbs(self, points, points_features, viewpoints, t):
        result_node_ids = query_grid(points, self.xyz_mins, self.xyz_maxs, self.sons, self.is_leaf_node)

        assert bool(torch.all(result_node_ids != -1)), "Grid not found"

        features = self.features[result_node_ids]
        return self.decoder(features, points_features, viewpoints, t)
        # return torch.cat([self.sh_features_dc[result_node_ids], self.sh_features_rest[result_node_ids]], dim=1)
    
    def save_weights(self, path):
        torch.save((self.features,
                    self.decoder.state_dict(),
                    self.sh_features_dc,
                    self.sh_features_rest,
                    self.xyz_mins,
                    self.xyz_maxs,
                    self.sons,
                    self.is_leaf_node), path)
    
    def load_weights(self, path):
        (self.features,
        decoder_state,
        self.sh_features_dc,
        self.sh_features_rest,
        self.xyz_mins,
        self.xyz_maxs,
        self.sons,
        self.is_leaf_node) = torch.load(path)

        self.decoder.load_state_dict(decoder_state)
        
            

        